#include <stdio.h>
#include <string.h>
#include <omp.h>
#include "lua.h"
#include "luaT.h"
#include "TH/TH.h"
#include "lauxlib.h"
#include "lualib.h"

#include "ransampl.h"

static int roulette_sampling(lua_State* L) {
	THDoubleTensor* data = luaT_toudata(L, -1, "torch.DoubleTensor");
	int free_data = 0;
	if (!THDoubleTensor_isContiguous(data)) {
		data = THDoubleTensor_newContiguous(data);
		free_data = 1;
	}
	double* probs = THDoubleTensor_data(data);
	int nrow = THDoubleTensor_size(data, 0);
	int ncol = THDoubleTensor_size(data, 1);
	
	THLongTensor* sample = THLongTensor_newWithSize1d(ncol);
	THDoubleTensor* rand = THDoubleTensor_newWithSize1d(ncol);
	THDoubleTensor_uniform(rand, 0, 1);
	
	int icol;
	#pragma omp parallel for shared(sample,rand,ncol,nrow,probs) private(icol)
	for (icol = 0; icol < ncol; icol++) {
		double r = THDoubleTensor_get1d(rand, icol);
		//printf("% f", r); 
		double acc = probs[icol];
		int irow;
		for (irow = 1; irow < nrow + 1; irow++) {
			if ( r <= acc ) {
				THLongTensor_set1d(sample, icol, irow);
				break;
			}
			acc += probs[irow*ncol+icol];
		}
	}

	if (free_data) 
		THDoubleTensor_free(data);
	THDoubleTensor_free(rand);
	luaT_pushudata(L, sample, "torch.LongTensor");
	return 1;
}

// indepedent chain metropolis-hastings
// states[:,i] is chain of states drawn from a transition distribution
// ratios[:,i] is unnormalized_true_prob / unnormalized_transition_prob
static void indep_chain_metropolis_hastings_(THLongTensor* states, 
											THDoubleTensor* ratios) {
	long chainlen = THLongTensor_size(states, 0);
	long nsamples = THLongTensor_size(states, 1);

	THDoubleTensor* r = THDoubleTensor_newWithSize2d(chainlen, nsamples);
	THDoubleTensor_uniform(r, 0, 1);

	int i;
	#pragma omp parallel for shared(states, ratios, r) private(i)
	for (i = 0; i < nsamples; i++) {
		int j;
		for (j = 1; j < chainlen; j++) {
			long cur_state = THLongTensor_get2d(states, j-1, i);
			double cur_ratio = THDoubleTensor_get2d(ratios, j-1, i);
			//long proposed_state = THLongTensor_get2d(states, j, i);
			double proposed_ratio = THDoubleTensor_get2d(ratios, j, i);

			double r_ = THDoubleTensor_get2d(r, j, i);
			double accept_rate = proposed_ratio / cur_ratio;
			if (accept_rate > 1) accept_rate = 1;

			if (r_ > accept_rate) {
				THLongTensor_set2d(states, j, i, cur_state);
				THDoubleTensor_set2d(ratios, j, i, cur_ratio);
			}
		} 
	}

	THDoubleTensor_free(r);
}

static int indep_chain_metropolis_hastings(lua_State* L) {
	THLongTensor* states = luaT_toudata(L, -2, "torch.LongTensor");
	THDoubleTensor* ratios = luaT_toudata(L, -1, "torch.DoubleTensor");
	indep_chain_metropolis_hastings_(states, ratios);
	return 0;
}

// alias sampling
// prob is a vector
static ransampl_ws* init_alias_sampling_(double* prob, int n) {
	ransampl_ws* ws = ransampl_alloc(n);
	ransampl_set(ws, prob);

	return ws;
}

static int init_alias_sampling(lua_State* L) {
	THDoubleTensor* prob = THDoubleTensor_newContiguous(
								luaT_toudata(L, -1, "torch.DoubleTensor"));
	double* prob__ = THDoubleTensor_data(prob);
	int n = THDoubleTensor_size(prob, 0);
	ransampl_ws* ws = init_alias_sampling_(prob__, n);
	
	THIntStorage* alias_ = THIntStorage_newWithData(ws->alias, ws->n);
	THDoubleStorage* prob_ = THDoubleStorage_newWithData(ws->prob, ws->n);
	
	luaT_pushudata(L, alias_, "torch.IntStorage");
	luaT_pushudata(L, prob_, "torch.DoubleStorage");

	free(ws);
	return 2;
}

static THLongTensor* gen_alias_sampling_(int* alias, double* prob, int n, int nsample) {
	ransampl_ws *ws = malloc( sizeof(ransampl_ws) );
	ws->n = n;
	ws->alias = alias;
	ws->prob = prob;

	THDoubleTensor* r = THDoubleTensor_newWithSize2d(2,nsample);
	THDoubleTensor_uniform(r, 0, 1);

	THLongTensor* sample = THLongTensor_newWithSize1d(nsample);
	int i;
	for (i = 0; i < nsample; i++) {
		double a = THDoubleTensor_get2d(r, 0, i);
		double b = THDoubleTensor_get2d(r, 1, i);
		int x = ransampl_draw( ws, a, b);
		THLongTensor_set1d(sample, i, x + 1); // note that counting starts at 1 (in Lua)
	}
	
	THDoubleTensor_free(r);
	free(ws);
	return sample;
}

static int gen_alias_sampling(lua_State* L) {
	THIntStorage* alias_ = luaT_toudata(L, -3, "torch.IntStorage");
	THDoubleStorage* prob_ = luaT_toudata(L, -2, "torch.DoubleStorage");
	int nsample = lua_tonumber(L, -1);
	int n = THIntStorage_size(alias_);

	THLongTensor* sample = gen_alias_sampling_(THIntStorage_data(alias_), 
											THDoubleStorage_data(prob_),
											n, nsample);

	luaT_pushudata(L, sample, "torch.LongTensor");

	return 1;
}

// if mv == false	u[i,:] = v[v_id[i],:] * M[M_id[i],:,:]
// if mv == true	u[i,:] = M[M_id[i],:,:] * v[v_id[i],:]
static THDoubleTensor* multi_matrix_mul_vector_(
						THDoubleTensor* M, THDoubleTensor* v, 
						int mv) {
	long n	= THDoubleTensor_size(M, 0);
	long n1	= THDoubleTensor_size(M, 1);
	long n2 = THDoubleTensor_size(M, 2);
	THDoubleTensor* u = mv ? 	THDoubleTensor_newWithSize2d(n, n1) : 
								THDoubleTensor_newWithSize2d(n, n2);

	int i;

	//omp_set_num_threads(16);
	//#pragma omp parallel for shared(M, v, u) private(i)
	for (i = 0; i < n; i++) {
		THDoubleTensor* M_ = THDoubleTensor_newSelect(M, 0, i);
		if (!mv) THDoubleTensor_transpose(M_, NULL, 0, 1);
        THDoubleTensor* v_ = THDoubleTensor_newSelect(v, 0, i);
        THDoubleTensor* u_ = THDoubleTensor_newSelect(u, 0, i);
		THDoubleTensor_addmv(u_, 0, u_, 1, M_, v_);
		THDoubleTensor_free(M_);
		THDoubleTensor_free(v_);
		THDoubleTensor_free(u_);
	}

	return u;
}

static int multi_vm(lua_State* L) {
	THDoubleTensor* v = luaT_toudata(L, 1, "torch.DoubleTensor");
	THDoubleTensor* M = luaT_toudata(L, 2, "torch.DoubleTensor");

	THDoubleTensor* u = multi_matrix_mul_vector_(M, v, 0);
	luaT_pushudata(L, u, "torch.DoubleTensor");
	return 1;
}

static int multi_mv(lua_State* L) {
	THDoubleTensor* M = luaT_toudata(L, 1, "torch.DoubleTensor");
	THDoubleTensor* v = luaT_toudata(L, 2, "torch.DoubleTensor");

	THDoubleTensor* u = multi_matrix_mul_vector_(M, v, 1);
	luaT_pushudata(L, u, "torch.DoubleTensor");
	return 1;
}

int luaopen_cutils(lua_State *L) {
	lua_register(L, "roulette_sampling", roulette_sampling);
	lua_register(L, "indep_chain_metropolis_hastings", indep_chain_metropolis_hastings);
	lua_register(L, "init_alias_sampling", init_alias_sampling);
	lua_register(L, "gen_alias_sampling", gen_alias_sampling);
	lua_register(L, "multi_vm", multi_vm);
	lua_register(L, "multi_mv", multi_mv);
	return 0;
}
